#include "yolo_byte_traker/yolo_byte_tracker.hpp"
#include <random>
cv::Scalar getRandomColor()
{
    static std::random_device rd;
    static std::mt19937 gen(rd());
    static std::uniform_int_distribution<int> dis(0, 255);

    int r = dis(gen);
    int g = dis(gen);
    int b = dis(gen);

    return cv::Scalar(b, g, r);
}

yolo_byte_traker::yolo_byte_traker():Node("yolo_byte_traker_node"){
        // -----------------------------------------declare parameters-----------------------------------------
        this->declare_parameter("subscribers.camera_reading.topic", "/camera/images");
        this->declare_parameter("subscribers.camera_reading.queue_size", 10);

        this->declare_parameter("publishers.detection_image.topic", "/detection_image/result_image");
        this->declare_parameter("publishers.detection_image.queue_size", 10);
        this->declare_parameter("publishers.detection_image.latch", true);

        this->declare_parameter("publishers.tracker_image.topic", "/tracker_image/result_image");
        this->declare_parameter("publishers.tracker_image.queue_size", 10);
        this->declare_parameter("publishers.tracker_image.latch", true);

        this->declare_parameter("publishers.traker_Box.topic", "/tracker_image/result_box");
        this->declare_parameter("publishers.traker_Box.queue_size", 10);
        this->declare_parameter("publishers.traker_Box.latch", true);


        this->declare_parameter("yolo_model.weight_file.name", "/home/nvidia/soft_simulate/src/yolo_byte_traker/source/best.onnx");
        this->declare_parameter("yolo_model.threshold.value", 0.5);
        this->declare_parameter("yolo_model.nmsthreshold.value", 0.45);

        this->declare_parameter("yolo_model.frame_rate.value", 30);
        this->declare_parameter("yolo_model.buf_size.value", 30);

        this->declare_parameter("yolo_model.topK.value", 100);
        this->declare_parameter("yolo_model.detection_classes.names", std::vector<std::string>());
        this->declare_parameter<std::string>("subscribers.call_server.topic","/drone_target_tracker_server_topic");
        this->declare_parameter<std::string>("services.calldetectserver.names","/yolo_byte_traker_server");
        
        this->declare_parameter<int>("parameters.vehicle_id",1);
        this->declare_parameter<std::string>("parameters.namesspace","drone_1");

        this->declare_parameter<std::string>("subscribers.global_pose.topic","/global_position/global");
        this->declare_parameter<std::string>("publishers.detection_label.topic","/tracker_image/result_label");
        
        // -----------------------------------------read paramer-----------------------------------------
        yolo_config.conf_threshold = 0.5;
        yolo_config.conf_threshold = this->get_parameter("yolo_model.threshold.value").as_double();

        yolo_config.model_path = "/home/nvidia/soft_simulate/src/yolo_byte_traker/source/best.engine";
        yolo_config.model_path = this->get_parameter("yolo_model.weight_file.name").as_string();

        yolo_config.class_names={};
        yolo_config.class_names = this->get_parameter("yolo_model.detection_classes.names").as_string_array();

        yolo_config.nms_threshold = 0.4;
        yolo_config.nms_threshold = this->get_parameter("yolo_model.nmsthreshold.value").as_double();

        namespaces = this->get_parameter("parameters.namesspace").as_string();
        RCLCPP_INFO(this->get_logger(),"namespaces is : %s",namespaces.c_str());

        int frame_rate = 30;
        int buffer_size = 30;

        frame_rate = this->get_parameter("yolo_model.frame_rate.value").as_int();
        buffer_size = this->get_parameter("yolo_model.buf_size.value").as_int();
        if (yolo_config.class_names.empty()|| yolo_config.model_path.empty() ||frame_rate == 0 || buffer_size == 0 )
        {
          rclcpp::shutdown();
        }
        byte_tracker.init_ByteTracker(frame_rate,buffer_size);
        detectionImageTopicName = "/detection_image/result_image";
        int detectionImageQueueSize = 10;
        bool detectionImageLatch = true;
        detectionImageTopicName = this->get_parameter("publishers.detection_image.topic").as_string();
        detectionImageTopicName = namespaces + detectionImageTopicName;
        RCLCPP_INFO(this->get_logger(),"detectionImageTopicName is : %s",detectionImageTopicName.c_str());
        detectionImageQueueSize = this->get_parameter("publishers.detection_image.queue_size").as_int();
        detectionImageLatch = this->get_parameter("publishers.detection_image.latch").as_bool();

        TrackerImageTopicName= "/detection_image/tracker_image";
        int TrackerImageQueueSize = 1;
        bool TrackerImageLatch = true;
        TrackerImageTopicName = this->get_parameter("publishers.tracker_image.topic").as_string();
        TrackerImageTopicName = namespaces + TrackerImageTopicName;
        TrackerImageQueueSize = this->get_parameter("publishers.tracker_image.queue_size").as_int();
        TrackerImageLatch = this->get_parameter("publishers.tracker_image.latch").as_bool();
        
        trakerBoxTopicName = "/detection_image/traker_box";

        int trakerBoxQueueSize = 1;
        bool trakerBoxLatch = true;
        trakerBoxTopicName = this->get_parameter("publishers.traker_Box.topic").as_string();
        trakerBoxTopicName = namespaces + trakerBoxTopicName;
        RCLCPP_INFO(this->get_logger(),"trakerBoxTopicName is : %s",trakerBoxTopicName.c_str());

        trakerBoxQueueSize = this->get_parameter("publishers.traker_Box.queue_size").as_int();
        trakerBoxLatch  = this->get_parameter("publishers.traker_Box.latch").as_bool();
        

        // int vehicle_id = 1;
        vehcile_id_ = this->get_parameter("parameters.vehicle_id").as_int();
        // service
        service_name = this->get_parameter("services.calldetectserver.names").as_string();
        service_name = namespaces + service_name;
        cameraTopicName = "camera/image_raw";
        int cameraQueueSize = 1;
        cameraTopicName = this->get_parameter("subscribers.camera_reading.topic").as_string();
        cameraQueueSize = this->get_parameter("subscribers.camera_reading.queue_size").as_int();

        gloabal_position_topic = this->get_parameter("subscribers.global_pose.topic").as_string();
        gloabal_position_topic = namespaces + gloabal_position_topic;
        location_trakerBoxresult_topic = this->get_parameter("publishers.detection_label.topic").as_string();
        location_trakerBoxresult_topic = namespaces + location_trakerBoxresult_topic;

        this->get_parameter("subscribers.call_server.topic",call_server_topic_name);
        call_server_topic_name = namespaces + call_server_topic_name;
        RCLCPP_INFO(this->get_logger(),"call_server_topic_name is : %s",call_server_topic_name.c_str());
        yolo_traker.init(yolo_config);
        service_status = true;
  
    }


yolo_byte_traker::~yolo_byte_traker() {
 RCLCPP_INFO(this->get_logger(),"[YoloObjectDetector] Node shutting down.");
}

bool yolo_byte_traker::readParameters() {
 RCLCPP_INFO(this->get_logger(),"[YoloObjectDetector] readparamers");
  
  yolo_config.conf_threshold = 0.5;
  yolo_config.conf_threshold = this->get_parameter("yolo_model.threshold.value").as_double();

  yolo_config.model_path = "/home/nvidia/soft_simulate/src/yolo_byte_traker/source/best.engine";
  yolo_config.model_path = this->get_parameter("yolo_model.weight_file.name").as_string();

  yolo_config.class_names={};
  yolo_config.class_names = this->get_parameter("yolo_model.detection_classes.names").as_string_array();

  yolo_config.nms_threshold = 0.4;
  yolo_config.nms_threshold = this->get_parameter("yolo_model.nmsthreshold.value").as_double();

  namespaces = this->get_parameter("parameters.namesspace").as_string();
  RCLCPP_INFO(this->get_logger(),"namespaces is : %s",namespaces.c_str());

  int frame_rate = 30;
  int buffer_size = 60;

  frame_rate = this->get_parameter("yolo_model.frame_rate.value").as_int();
  buffer_size = this->get_parameter("yolo_model.buf_size.value").as_int();
 

  if (yolo_config.class_names.empty()|| yolo_config.model_path.empty() ||frame_rate == 0 || buffer_size == 0 )
  {
    return false;
  }
  byte_tracker.init_ByteTracker(frame_rate,buffer_size);
  return true;
}

void yolo_byte_traker::init() {

  image_transport::ImageTransport it(shared_from_this());
  // imageSubscriber_ = it.subscribe(cameraTopicName, 10, std::bind(&yolo_byte_traker::cameraCallback, this, std::placeholders::_1));
  cameraSubscriber_ = this->create_subscription<sensor_msgs::msg::Image>(cameraTopicName,rclcpp::QoS(10).best_effort(),std::bind(&yolo_byte_traker::cameraCallback, this, std::placeholders::_1));
  detectionImagePublisher_ = this->create_publisher<sensor_msgs::msg::Image>(detectionImageTopicName, 10);
  // detectionImagePublisher_ = nodeHandle_.advertise<sensor_msgs::Image>(detectionImageTopicName, detectionImageQueueSize, detectionImageLatch);
  
  trakerImagePublisher_ = this->create_publisher<sensor_msgs::msg::Image>(TrackerImageTopicName, 10);
  // trakerImagePublisher_ = nodeHandle_.advertise<sensor_msgs::Image>(TrackerImageTopicName, TrackerImageQueueSize, TrackerImageLatch);
  trakerBoxresult = this->create_publisher<yolo_detect_message::msg::TrakerBoxs>(trakerBoxTopicName, 10);
  // trakerBoxresult = nodeHandle_.advertise<yolo_detect_message::TrakerBoxs>(trakerBoxTopicName, trakerBoxQueueSize, trakerBoxLatch);
  
  callback_group_ = this->create_callback_group(rclcpp::CallbackGroupType::MutuallyExclusive);
  service_ = this->create_service<yolo_detect_message::srv::Service>(service_name, std::bind(&yolo_byte_traker::service_callback,this,std::placeholders::_1,std::placeholders::_2),rmw_qos_profile_services_default,callback_group_);
  service_status_subscriber_ = this->create_subscription<yolo_detect_message::msg::CallServerStatus>(call_server_topic_name,rclcpp::QoS(5).best_effort(),std::bind(&yolo_byte_traker::service_status_callback,this,std::placeholders::_1));
  global_subscription_ = this->create_subscription<sensor_msgs::msg::NavSatFix>(
            gloabal_position_topic, 10,
            std::bind(&yolo_byte_traker::uav_global_position_callback, this, std::placeholders::_1));
  location_trakerBoxresult = this->create_publisher<yolo_detect_message::msg::LocationTrakerBoxs>(location_trakerBoxresult_topic, 10);
  info_topic();
}

void yolo_byte_traker::cameraCallback(const std::shared_ptr<const sensor_msgs::msg::Image> msg) {
  if(msg == nullptr) 
  {
    return;
  }
  
  cv::Mat image = cv_bridge::toCvCopy(msg, "bgr8")->image;
  if(image.empty()){
    RCLCPP_DEBUG(this->get_logger(),"[yolo byte tracker] image is empty");
    return;
  }

  cv::Mat camImageCopy_ = image.clone();
  auto result = yolo_traker.detect(camImageCopy_);
  vector<Object> temp{};
  if(result.empty()) {
    RCLCPP_INFO(this->get_logger(),"[yolo byte tracker] no object detected");
    image_tag++;
    // std::string image_name = image_path + to_string(image_tag) + ".png";
    // cv::imwrite(image_name, image);

  }
  
  for (int i = 0; i < result.size(); i++)
  {
      Object value;
      value.classLab = result.at(i).lable_name;
      value.label =result.at(i).label;
      value.prob = result.at(i).score;
      value.rect = cv::Rect(result.at(i).x1,result.at(i).x2,result.at(i).width,result.at(i).height);
      temp.push_back(value);
  }
  int w = image.cols;
  int h = image.rows;

  auto traker_result = byte_tracker.update(temp);
  yolo_detect_message::msg::TrakerBoxs traker_box;
  yolo_detect_message::msg::TrakerBox box;
  
  yolo_detect_message::msg::LocationTrakerBox location_traker_box;
  yolo_detect_message::msg::LocationTrakerBoxs location_traker_boxs;
  for (size_t i = 0; i < traker_result.size(); i++)
  {
    box.trakerid = traker_result[i].track_id;
  
    box.xmin = traker_result[i].tlwh[0];
    box.ymin = traker_result[i].tlwh[1];
    box.width  = traker_result[i].tlwh[2];
    box.height  = traker_result[i].tlwh[3];
    box.probability = traker_result[i].score;
    box.class_label = traker_result[i].classLabel;
 
    box.id = traker_result[i].classId;
    traker_box.traker_boxes.push_back(box);
  }
  for (size_t i = 0; i < traker_result.size(); i++)
  {
    //correct on 2025-01-21
    if(traker_result[i].score >= 85.0)
    {
      location_traker_box.trakerid = traker_result[i].track_id;
  
      location_traker_box.xmin = traker_result[i].tlwh[0];
      location_traker_box.ymin = traker_result[i].tlwh[1];
      location_traker_box.width  = traker_result[i].tlwh[2];
      location_traker_box.heigh  = traker_result[i].tlwh[3];
      location_traker_box.probability = traker_result[i].score;
      location_traker_box.class_label = traker_result[i].classLabel;
    
      location_traker_box.id = traker_result[i].classId;

      int err = 20;
      int err_x = (location_traker_box.xmin + traker_result[i].tlwh[2]) / 2;
      int err_y = (location_traker_box.ymin + traker_result[i].tlwh[3]) / 2;
      if ((err_x > -1 * err && err_x <  err)&& (err_y > -1 * err && err_y <  err)){
        location_traker_box.is_false = true;
        location_traker_box.altitude = uav_gloabal_position.altitude;
        location_traker_box.longtitude = uav_gloabal_position.longitude;
        location_traker_box.latitude = uav_gloabal_position.latitude;
      }else{
        location_traker_box.is_false = false;
        location_traker_box.altitude = 0;
        location_traker_box.longtitude = 0;
        location_traker_box.latitude = 0;
      }
    }
    location_traker_boxs.traker_boxes.push_back(location_traker_box);
  }
  cv::Mat image_result1  = Draw_Line(camImageCopy_,traker_result);
  cv::Mat image_result = yolo_traker.YOLOv10_Draw_Line(result,camImageCopy_);
  publishDetectionImage(image_result);
  publishTrackerImage(image_result1);
  std_msgs::msg::Header header;
  traker_box.header.stamp = this->get_clock()->now();
  trakerBoxresult->publish(traker_box);
  location_traker_boxs.header.stamp = this->get_clock()->now();
  location_trakerBoxresult->publish(location_traker_boxs);

  return;
}


bool yolo_byte_traker::publishDetectionImage(const cv::Mat& detectionImage) {
  // if (detectionImagePublisher_.getNumSubscribers() < 1) return false;
  cv_bridge::CvImage cvImage;

  std_msgs::msg::Header header;
  header.stamp = this->now();
  header.frame_id = "camera_frame";
  
  auto img_msg = cv_bridge::CvImage(header, "bgr8", detectionImage).toImageMsg();
  // cvImage.header.stamp = ros::Time::now();
  // cvImage.header.frame_id = "detection_image";
  // cvImage.encoding = sensor_msgs::image_encodings::BGR8;
  // cvImage.image = detectionImage.clone(); 
  detectionImagePublisher_->publish(*img_msg);
  RCLCPP_DEBUG(this->get_logger(),"Detection image has been published.");
  return true;
}

cv::Mat yolo_byte_traker::Draw_Line(const cv::Mat &image, const std::vector<STrack> &detections)
{
  if(image.empty() || detections.size() == 0)
  {
    return;
  }

  cv::Mat result = image.clone();
  for (const auto &detection : detections)
  {
      // 绘制边界框，颜色为绿色（BGR: 0, 255, 0），线宽为2
      cv::Scalar color = getRandomColor();
      cv::Rect rect = {detection.tlwh[0],detection.tlwh[1],detection.tlwh[2],detection.tlwh[3]};
      cv::rectangle(result, rect, color, 3);
      std::string labels = format("%.2f", detection.score);
      // 准备标签文本，包含类别名称和置信度
      std::string label = detection.classLabel + ": " + labels + " id:" + to_string(detection.track_id);

      // 计算标签文本的尺寸和基线
      int baseLine;
      cv::Size labelSize = cv::getTextSize(label, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);

      // 绘制标签背景矩形，颜色为白色（BGR: 255, 255, 255）
      cv::rectangle(
          result,
          cv::Point(rect.x, rect.y - labelSize.height),
          cv::Point(rect.x + labelSize.width, rect.y + baseLine),
          color,
          cv::FILLED);

      // 绘制标签文本，颜色为黑色（BGR: 0, 0, 0）
      cv::putText(
          result,
          label,
          cv::Point(rect.x, rect.y),
          cv::FONT_HERSHEY_SIMPLEX,
          0.5,
          cv::Scalar(255, 0, 0),
          1);
  }

  return result;
}

bool yolo_byte_traker::publishTrackerImage(const cv::Mat& detectionImage) {
  // if (detectionImagePublisher_.getNumSubscribers() < 1) return false;
  // cv_bridge::CvImage cvImage;
  // cvImage.header.stamp = ros::Time::now();
  // cvImage.header.frame_id = "detection_image";
  // cvImage.encoding = sensor_msgs::image_encodings::BGR8;
  // cvImage.image = detectionImage.clone(); 
  // trakerImagePublisher_.publish(*cvImage.toImageMsg());

  cv_bridge::CvImage cvImage;
  std_msgs::msg::Header header;
  header.stamp = this->now();
  header.frame_id = "camera_frame";
  
  auto img_msg = cv_bridge::CvImage(header, "bgr8", detectionImage).toImageMsg();
  RCLCPP_DEBUG(this->get_logger(),"Detection image has been published.");
  trakerImagePublisher_->publish(*img_msg);
  return true;
}


void yolo_byte_traker::service_callback(const yolo_detect_message::srv::Service::Request::SharedPtr req,
                        const yolo_detect_message::srv::Service::Response::SharedPtr res){
    bool reqs = req->status;
    if(reqs){
        RCLCPP_INFO(this->get_logger(),"[YoloObjectDetector] service callback.");
    }
    else{
        RCLCPP_INFO(this->get_logger(),"[YoloObjectDetector] service callback false.");

    }
    {
        std::lock_guard<std::mutex> lock_image(service_status_mutex);
        service_status = reqs;
    }
    
    res->success = true;
    res->drone_statue = "yolo byte tracker service_status is " + std::to_string(service_status);
}



void yolo_byte_traker::service_status_callback(const yolo_detect_message::msg::CallServerStatus::SharedPtr msg){
  
  bool reqs = msg->status;
    if(reqs != service_status)
    {
        std::lock_guard<std::mutex> lock_image(service_status_mutex);
        service_status = reqs;
    }
    RCLCPP_INFO(this->get_logger(),"[yolo byte tracker] service status is %d.",service_status);
}



void yolo_byte_traker::info_topic(){
  RCLCPP_INFO(this->get_logger(),"namespaces: %s",namespaces.c_str());
  RCLCPP_INFO(this->get_logger(),"[YoloObjectDetector] info topic. start.");
  RCLCPP_INFO(this->get_logger(), "vehicle id: %d", vehcile_id_);
  RCLCPP_INFO(this->get_logger(),"detectionImageTopicName: %s",detectionImageTopicName.c_str());
  RCLCPP_INFO(this->get_logger(),"cameraTopicName: %s",cameraTopicName.c_str());
  RCLCPP_INFO(this->get_logger(),"TrackerImageTopicName: %s",TrackerImageTopicName.c_str());
  RCLCPP_INFO(this->get_logger(),"trakerBoxTopicName: %s",trakerBoxTopicName.c_str());
  RCLCPP_INFO(this->get_logger(),"service_name: %s",service_name.c_str());
  
  RCLCPP_INFO(this->get_logger(),"call_server_topic_name: %s",call_server_topic_name.c_str());
  RCLCPP_INFO(this->get_logger(),"gloabal_position_topic: %s",gloabal_position_topic.c_str());
 
  RCLCPP_INFO(this->get_logger(),"location_trakerBoxresult_topic: %s",location_trakerBoxresult_topic.c_str());
        
  RCLCPP_INFO(this->get_logger(),"[YoloObjectDetector] info topic. edd.");
}
void yolo_byte_traker::uav_global_position_callback(const sensor_msgs::msg::NavSatFix::SharedPtr  msg){
   uav_gloabal_position = *msg;
}